{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "collapsed": true,
        "id": "fHaTpMLdneR2"
      },
      "source": [
        "### Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "wHdNs60yneR5"
      },
      "outputs": [],
      "source": [
        "# -*- coding: utf-8 -*-\n",
        "\"\"\"\n",
        "lstm_text_generator_word_level.py\n",
        "\n",
        "该模块实现了一个基于 NumPy 的词级别文本生成器，使用 LSTM 模型。\n",
        "代码包含训练和文本生成功能，适用于教学和理解 LSTM 的工作原理。\n",
        "\n",
        "Created on Mon Nov 25 11:17:34 2019\n",
        "@author: lizhenping\n",
        "\"\"\"\n",
        "\n",
        "import numpy as np\n",
        "import re\n",
        "from collections import Counter\n",
        "\n",
        "class Tokenizer:\n",
        "    \"\"\"\n",
        "    词级别的分词器，负责将文本转换为词语索引序列，以及索引序列转换为文本。\n",
        "    \"\"\"\n",
        "    def __init__(self, text, max_vocab_size=None):\n",
        "        self.special_tokens = ['<PAD>', '<UNK>', '<EOS>']  # 特殊标记\n",
        "        self.max_vocab_size = max_vocab_size\n",
        "        self.build_vocab(text)\n",
        "\n",
        "    def build_vocab(self, text):\n",
        "        # 使用正则表达式分词，保留标点符号\n",
        "        words = re.findall(r'\\w+|[^\\s\\w]+', text)\n",
        "        word_counts = Counter(words)\n",
        "\n",
        "        if self.max_vocab_size:\n",
        "            # 根据词频排序，取前 max_vocab_size 个词\n",
        "            most_common = word_counts.most_common(self.max_vocab_size - len(self.special_tokens))\n",
        "            vocab = self.special_tokens + [word for word, _ in most_common]\n",
        "        else:\n",
        "            vocab = self.special_tokens + list(word_counts.keys())\n",
        "\n",
        "        self.word_to_ix = {word: i for i, word in enumerate(vocab)}\n",
        "        self.ix_to_word = {i: word for i, word in enumerate(vocab)}\n",
        "        self.vocab_size = len(vocab)\n",
        "\n",
        "    def text_to_sequence(self, text):\n",
        "        # 将文本转换为词语索引序列\n",
        "        words = re.findall(r'\\w+|[^\\s\\w]+', text)\n",
        "        sequence = []\n",
        "        for word in words:\n",
        "            sequence.append(self.word_to_ix.get(word, self.word_to_ix['<UNK>']))\n",
        "        sequence.append(self.word_to_ix['<EOS>'])\n",
        "        return sequence\n",
        "\n",
        "    def sequence_to_text(self, sequence):\n",
        "        # 将词语索引序列转换为文本\n",
        "        words = [self.ix_to_word.get(idx, '<UNK>') for idx in sequence]\n",
        "        text = ''.join(words)  # 对于中文不需要空格分隔\n",
        "        return text\n",
        "\n",
        "class Module:\n",
        "    \"\"\"\n",
        "    模块基类，所有神经网络模块的父类。\n",
        "    提供参数管理、梯度清零、参数保存和加载功能。\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        self.parameters = []\n",
        "        self.gradients = []\n",
        "\n",
        "    def zero_grad(self):\n",
        "        \"\"\"\n",
        "        将所有参数的梯度清零。\n",
        "        \"\"\"\n",
        "        for grad in self.gradients:\n",
        "            grad.fill(0)\n",
        "\n",
        "    def init_weights(self, init_range):\n",
        "        \"\"\"\n",
        "        初始化模型权重，服从[-init_range, init_range]的均匀分布。\n",
        "        \"\"\"\n",
        "        for param in self.parameters:\n",
        "            param[:] = np.random.uniform(-init_range, init_range, size=param.shape)\n",
        "\n",
        "class LSTMCell(Module):\n",
        "    \"\"\"\n",
        "    实现单个 LSTM 单元，包括前向和反向传播。\n",
        "    \"\"\"\n",
        "    def __init__(self, input_size, hidden_size, init_range=0.1):\n",
        "        super().__init__()\n",
        "        self.input_size = input_size\n",
        "        self.hidden_size = hidden_size\n",
        "\n",
        "        # 初始化权重和偏置\n",
        "        self.W_ih = np.random.uniform(-init_range, init_range, (4 * hidden_size, input_size))\n",
        "        self.W_hh = np.random.uniform(-init_range, init_range, (4 * hidden_size, hidden_size))\n",
        "        self.b_ih = np.zeros((4 * hidden_size, 1))\n",
        "        self.b_hh = np.zeros((4 * hidden_size, 1))\n",
        "\n",
        "        # 存储参数和对应的梯度\n",
        "        self.parameters = [self.W_ih, self.W_hh, self.b_ih, self.b_hh]\n",
        "        self.gradients = [np.zeros_like(param) for param in self.parameters]\n",
        "\n",
        "    def forward(self, x, h_prev, c_prev):\n",
        "        \"\"\"\n",
        "        前向传播\n",
        "\n",
        "        参数:\n",
        "        - x: 当前时间步的输入，形状为 (input_size, batch_size)\n",
        "        - h_prev: 前一时间步的隐藏状态，形状为 (hidden_size, batch_size)\n",
        "        - c_prev: 前一时间步的细胞状态，形状为 (hidden_size, batch_size)\n",
        "\n",
        "        返回:\n",
        "        - h_next: 当前时间步的隐藏状态\n",
        "        - c_next: 当前时间步的细胞状态\n",
        "        \"\"\"\n",
        "        self.x = x\n",
        "        self.h_prev = h_prev\n",
        "        self.c_prev = c_prev\n",
        "\n",
        "        # 计算门和候选值\n",
        "        gates = np.dot(self.W_ih, x) + self.b_ih + np.dot(self.W_hh, h_prev) + self.b_hh\n",
        "        self.i_gate = sigmoid(gates[0:self.hidden_size, :])\n",
        "        self.f_gate = sigmoid(gates[self.hidden_size:2*self.hidden_size, :])  \n",
        "        self.o_gate = sigmoid(gates[2*self.hidden_size:3*self.hidden_size, :])\n",
        "        self.g_gate = tanh(gates[3*self.hidden_size:4*self.hidden_size, :])\n",
        "\n",
        "        # 计算细胞状态和隐藏状态\n",
        "        self.c_next = self.f_gate * c_prev + self.i_gate * self.g_gate\n",
        "        self.h_next = self.o_gate * tanh(self.c_next)\n",
        "\n",
        "        return self.h_next, self.c_next\n",
        "\n",
        "    def backward(self, dh_next, dc_next, grad_clip=5):\n",
        "        \"\"\"\n",
        "        反向传播\n",
        "\n",
        "        参数:\n",
        "        - dh_next: 当前时间步的隐藏状态梯度，形状为 (hidden_size, batch_size)\n",
        "        - dc_next: 当前时间步的细胞状态梯度，形状为 (hidden_size, batch_size)\n",
        "        - grad_clip: 梯度裁剪阈值，默认为 5\n",
        "\n",
        "        返回:\n",
        "        - dx: 对输入 x 的梯度\n",
        "        - dh_prev: 对前一隐藏状态 h_prev 的梯度\n",
        "        - dc_prev: 对前一细胞状态 c_prev 的梯度\n",
        "        \"\"\"\n",
        "        # 计算梯度\n",
        "        do = dh_next * tanh(self.c_next) * dsigmoid(self.o_gate)\n",
        "        dc = dc_next + dh_next * self.o_gate * dtanh(tanh(self.c_next))\n",
        "        di = dc * self.g_gate * dsigmoid(self.i_gate)\n",
        "        df = dc * self.c_prev * dsigmoid(self.f_gate)\n",
        "        dg = dc * self.i_gate * dtanh(self.g_gate)\n",
        "\n",
        "        # 拼接各个门的梯度\n",
        "        d_gates = np.vstack((di, df, do, dg))\n",
        "\n",
        "        # 计算权重和偏置的梯度\n",
        "        dW_ih = np.dot(d_gates, self.x.T)\n",
        "        dW_hh = np.dot(d_gates, self.h_prev.T)\n",
        "        db_ih = np.sum(d_gates, axis=1, keepdims=True)\n",
        "        db_hh = np.sum(d_gates, axis=1, keepdims=True)\n",
        "\n",
        "        # 裁剪梯度，避免梯度爆炸\n",
        "        dW_ih = np.clip(dW_ih, -grad_clip, grad_clip)\n",
        "        dW_hh = np.clip(dW_hh, -grad_clip, grad_clip)\n",
        "        db_ih = np.clip(db_ih, -grad_clip, grad_clip)\n",
        "        db_hh = np.clip(db_hh, -grad_clip, grad_clip)\n",
        "\n",
        "        # 计算输入、前一隐藏状态和前一细胞状态的梯度\n",
        "        dx = np.dot(self.W_ih.T, d_gates)\n",
        "        dh_prev = np.dot(self.W_hh.T, d_gates)\n",
        "        dc_prev = dc * self.f_gate\n",
        "\n",
        "        # 存储梯度\n",
        "        self.gradients[0] += dW_ih  \n",
        "        self.gradients[1] += dW_hh\n",
        "        self.gradients[2] += db_ih\n",
        "        self.gradients[3] += db_hh\n",
        "\n",
        "        return dx, dh_prev, dc_prev\n",
        "\n",
        "def sigmoid(x):\n",
        "    return 1 / (1 + np.exp(-x))\n",
        "\n",
        "def dsigmoid(y):\n",
        "    # y = sigmoid(x)\n",
        "    return y * (1 - y)\n",
        "\n",
        "def tanh(x):\n",
        "    return np.tanh(x)\n",
        "\n",
        "def dtanh(y):\n",
        "    # y = tanh(x)\n",
        "    return 1 - y ** 2\n",
        "\n",
        "class LSTMLayer(Module):\n",
        "    \"\"\"\n",
        "    LSTM 层，由多个 LSTM 单元组成。\n",
        "    \"\"\"\n",
        "    def __init__(self, input_size, hidden_size, num_layers=1, init_range=0.1):\n",
        "        super().__init__()\n",
        "        self.input_size = input_size\n",
        "        self.hidden_size = hidden_size\n",
        "        self.num_layers = num_layers\n",
        "        self.lstm_cells = [LSTMCell(input_size if i == 0 else hidden_size, hidden_size, init_range) \n",
        "                           for i in range(num_layers)]\n",
        "\n",
        "        # 存储参数和梯度\n",
        "        self.parameters = []\n",
        "        self.gradients = []\n",
        "        for lstm_cell in self.lstm_cells:\n",
        "            self.parameters.extend(lstm_cell.parameters)\n",
        "            self.gradients.extend(lstm_cell.gradients)\n",
        "\n",
        "    def forward(self, x, h0, c0):\n",
        "        \"\"\"\n",
        "        前向传播\n",
        "\n",
        "        参数:\n",
        "        - x: 输入序列，形状为 (input_size, seq_len, batch_size)\n",
        "        - h0: 初始隐藏状态列表，长度为 num_layers，每个元素形状为 (hidden_size, batch_size)\n",
        "        - c0: 初始细胞状态列表，长度为 num_layers，每个元素形状为 (hidden_size, batch_size)\n",
        "\n",
        "        返回:\n",
        "        - outputs: 每个时间步的输出，形状为 (hidden_size, seq_len, batch_size)\n",
        "        - (h_n, c_n): 最后一个时间步的隐藏状态和细胞状态列表\n",
        "        \"\"\"\n",
        "        seq_len = x.shape[1]\n",
        "        batch_size = x.shape[2]\n",
        "        h = [h0_layer.copy() for h0_layer in h0]\n",
        "        c = [c0_layer.copy() for c0_layer in c0]\n",
        "        outputs = []\n",
        "\n",
        "        for t in range(seq_len):\n",
        "            xt = x[:, t, :]  # 当前时间步的输入，形状 (input_size, batch_size)\n",
        "            for i, lstm_cell in enumerate(self.lstm_cells):\n",
        "                h_prev, c_prev = h[i], c[i]\n",
        "                h_next, c_next = lstm_cell.forward(xt, h_prev, c_prev)\n",
        "                h[i], c[i] = h_next, c_next\n",
        "                xt = h_next  # 当前层的输出作为下一层的输入\n",
        "            outputs.append(h[-1].reshape(self.hidden_size, 1, batch_size))  # 收集最后一层的输出\n",
        "\n",
        "        outputs = np.concatenate(outputs, axis=1)  # (hidden_size, seq_len, batch_size)\n",
        "        return outputs, (h, c)\n",
        "\n",
        "    def backward(self, doutputs, dh_n, dc_n, grad_clip=5): \n",
        "        \"\"\"\n",
        "        反向传播\n",
        "\n",
        "        参数:\n",
        "        - doutputs: 每个时间步的输出梯度，形状为 (hidden_size, seq_len, batch_size)\n",
        "        - dh_n: 最后一个时间步的隐藏状态梯度列表\n",
        "        - dc_n: 最后一个时间步的细胞状态梯度列表\n",
        "        - grad_clip: 梯度裁剪阈值，默认为 5\n",
        "\n",
        "        返回:\n",
        "        - dx: 输入序列的梯度，形状为 (input_size, seq_len, batch_size)\n",
        "        \"\"\"\n",
        "        seq_len = doutputs.shape[1]\n",
        "        batch_size = doutputs.shape[2]\n",
        "        dh_next = [dh.copy() for dh in dh_n]\n",
        "        dc_next = [dc.copy() for dc in dc_n]\n",
        "        dx = []\n",
        "\n",
        "        for t in reversed(range(seq_len)):\n",
        "            dh = doutputs[:, t, :]  # (hidden_size, batch_size)\n",
        "            for i in reversed(range(self.num_layers)):\n",
        "                lstm_cell = self.lstm_cells[i]\n",
        "                dh = dh + dh_next[i]\n",
        "                dx_step, dh_prev, dc_prev = lstm_cell.backward(dh, dc_next[i], grad_clip)\n",
        "                dh_next[i] = dh_prev\n",
        "                dc_next[i] = dc_prev\n",
        "                dh = dx_step  # 当前层的 dx_step 作为上一层的 dh\n",
        "            dx.insert(0, dx_step.reshape(self.input_size, 1, batch_size))\n",
        "\n",
        "        dx = np.concatenate(dx, axis=1)  # 拼接各时间步的输入梯度，形状 (input_size, seq_len, batch_size)\n",
        "        return dx  # 保持原形状返回\n",
        "\n",
        "class Embedding(Module):\n",
        "    \"\"\"\n",
        "    嵌入层，将词语索引转换为词向量。\n",
        "    \"\"\"\n",
        "    def __init__(self, vocab_size, embedding_dim, init_range=0.1):\n",
        "        super().__init__()\n",
        "        self.vocab_size = vocab_size\n",
        "        self.embedding_dim = embedding_dim\n",
        "\n",
        "        # 初始化嵌入矩阵\n",
        "        self.embedding_matrix = np.random.uniform(-init_range, init_range, (vocab_size, embedding_dim))\n",
        "\n",
        "        # 存储参数和梯度\n",
        "        self.parameters = [self.embedding_matrix]\n",
        "        self.gradients = [np.zeros_like(self.embedding_matrix)]\n",
        "\n",
        "    def forward(self, inputs):\n",
        "        \"\"\"\n",
        "        前向传播\n",
        "\n",
        "        参数:\n",
        "        - inputs: 词语索引序列，形状为 (seq_len, batch_size)\n",
        "\n",
        "        返回:\n",
        "        - outputs: 词向量序列，形状为 (embedding_dim, seq_len, batch_size)\n",
        "        \"\"\"\n",
        "        self.inputs = inputs\n",
        "        outputs = self.embedding_matrix[inputs]  # shape: (seq_len, batch_size, embedding_dim)\n",
        "        outputs = outputs.transpose(2, 0, 1)  # 转换为 (embedding_dim, seq_len, batch_size)\n",
        "        return outputs\n",
        "\n",
        "    def backward(self, doutputs):\n",
        "        \"\"\"\n",
        "        反向传播\n",
        "\n",
        "        参数:\n",
        "        - doutputs: 词向量序列的梯度，形状为 (embedding_dim, seq_len, batch_size)\n",
        "\n",
        "        返回:\n",
        "        - 无需返回输入梯度\n",
        "        \"\"\"\n",
        "        doutputs = doutputs.transpose(1, 2, 0)  # 转换为 (seq_len, batch_size, embedding_dim)\n",
        "        for i in range(self.inputs.shape[0]):  # seq_len\n",
        "            for j in range(self.inputs.shape[1]):  # batch_size\n",
        "                idx = self.inputs[i, j]\n",
        "                self.gradients[0][idx] += doutputs[i, j]\n",
        "        # 返回的梯度不传递给前面的层\n",
        "        return None\n",
        "\n",
        "class Linear(Module):\n",
        "    \"\"\"\n",
        "    线性层，实现 y = Wx + b。\n",
        "    \"\"\"\n",
        "    def __init__(self, input_size, output_size, bias=True, init_range=0.1):\n",
        "        super().__init__()\n",
        "        self.input_size = input_size\n",
        "        self.output_size = output_size\n",
        "        self.use_bias = bias\n",
        "\n",
        "        # 初始化权重和偏置\n",
        "        self.W = np.random.uniform(-init_range, init_range, (output_size, input_size))\n",
        "        self.b = np.zeros((output_size, 1)) if self.use_bias else None\n",
        "\n",
        "        # 存储参数和梯度\n",
        "        self.parameters = [self.W] if not self.use_bias else [self.W, self.b] \n",
        "        self.gradients = [np.zeros_like(param) for param in self.parameters]\n",
        "\n",
        "    def forward(self, inputs):\n",
        "        \"\"\"\n",
        "        前向传播\n",
        "\n",
        "        参数:\n",
        "        - inputs: 输入，形状为 (input_size, batch_size)\n",
        "\n",
        "        返回:\n",
        "        - outputs: 输出，形状为 (output_size, batch_size)\n",
        "        \"\"\"\n",
        "        self.inputs = inputs\n",
        "        outputs = np.dot(self.W, inputs)\n",
        "        if self.use_bias:\n",
        "            outputs += self.b\n",
        "        return outputs\n",
        "\n",
        "    def backward(self, doutputs):\n",
        "        \"\"\"\n",
        "        反向传播\n",
        "\n",
        "        参数:\n",
        "        - doutputs: 输出的梯度，形状为 (output_size, batch_size)\n",
        "\n",
        "        返回:\n",
        "        - dinputs: 输入的梯度，形状为 (input_size, batch_size)\n",
        "        \"\"\"\n",
        "        dinputs = np.dot(self.W.T, doutputs)\n",
        "        dW = np.dot(doutputs, self.inputs.T)\n",
        "        self.gradients[0] += dW\n",
        "        if self.use_bias:\n",
        "            db = np.sum(doutputs, axis=1, keepdims=True)\n",
        "            self.gradients[1] += db\n",
        "        return dinputs\n",
        "\n",
        "class CrossEntropyLoss:\n",
        "    \"\"\"\n",
        "    交叉熵损失函数\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        pass\n",
        "\n",
        "    def forward(self, inputs, targets, reduction='mean'):\n",
        "        \"\"\"\n",
        "        前向计算损失值\n",
        "\n",
        "        参数:\n",
        "        - inputs: 模型的输出，形状为 (num_classes, batch_size)\n",
        "        - targets: 目标类别索引，形状为 (batch_size,)\n",
        "        - reduction: 损失归约方式，可选 'mean' 或 'sum'，默认为 'mean'\n",
        "\n",
        "        返回:\n",
        "        - loss: 标量，平均或求和后的损失值\n",
        "        \"\"\"\n",
        "        self.inputs = inputs\n",
        "        self.targets = targets\n",
        "        num_classes, batch_size = inputs.shape\n",
        "\n",
        "        # 计算每个样本的对数 softmax\n",
        "        shifted_logits = inputs - np.max(inputs, axis=0, keepdims=True)\n",
        "        log_probs = shifted_logits - np.log(np.sum(np.exp(shifted_logits), axis=0, keepdims=True))\n",
        "        self.probs = np.exp(log_probs)\n",
        "        loss = -log_probs[self.targets, range(batch_size)]\n",
        "        if reduction == 'mean':\n",
        "            return np.mean(loss)\n",
        "        elif reduction == 'sum':\n",
        "            return np.sum(loss)\n",
        "\n",
        "    def backward(self):\n",
        "        \"\"\"\n",
        "        反向传播，计算输入的梯度\n",
        "\n",
        "        返回:\n",
        "        - dinputs: 模型输出的梯度，形状与 inputs 相同\n",
        "        \"\"\"\n",
        "        batch_size = self.inputs.shape[1]\n",
        "        dinputs = self.probs.copy()\n",
        "        dinputs[self.targets, range(batch_size)] -= 1\n",
        "        dinputs /= batch_size\n",
        "        return dinputs\n",
        "\n",
        "class AdamOptimizer:\n",
        "    \"\"\"\n",
        "    Adam 优化器\n",
        "    \"\"\"\n",
        "    def __init__(self, parameters, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):\n",
        "        self.parameters = parameters\n",
        "        self.lr = lr\n",
        "        self.beta1, self.beta2 = betas\n",
        "        self.eps = eps\n",
        "        self.weight_decay = weight_decay\n",
        "\n",
        "        self.m = [np.zeros_like(param) for param in parameters]  # 一阶矩\n",
        "        self.v = [np.zeros_like(param) for param in parameters]  # 二阶矩\n",
        "        self.t = 0  # 时间步\n",
        "\n",
        "    def step(self, gradients):\n",
        "        \"\"\"\n",
        "        应用梯度，更新参数，并更新时间步长。\n",
        "\n",
        "        参数:\n",
        "        - gradients: 与 parameters 对应的梯度列表\n",
        "        \"\"\"\n",
        "        self.t += 1  # 在参数更新之前递增时间步长\n",
        "\n",
        "        for idx, (param, grad) in enumerate(zip(self.parameters, gradients)):\n",
        "            # 应用权重衰减\n",
        "            if self.weight_decay != 0:\n",
        "                grad += self.weight_decay * param\n",
        "\n",
        "            # 更新一阶矩和二阶矩估计\n",
        "            self.m[idx] = self.beta1 * self.m[idx] + (1 - self.beta1) * grad\n",
        "            self.v[idx] = self.beta2 * self.v[idx] + (1 - self.beta2) * (grad ** 2)\n",
        "\n",
        "            # 修正一阶矩和二阶矩的偏差\n",
        "            m_hat = self.m[idx] / (1 - self.beta1 ** self.t)\n",
        "            v_hat = self.v[idx] / (1 - self.beta2 ** self.t)\n",
        "\n",
        "            # 更新参数\n",
        "            param -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)\n",
        "\n",
        "class LanguageModel(Module):\n",
        "    \"\"\"\n",
        "    语言模型，包括词嵌入层、LSTM 层和线性输出层。\n",
        "    \"\"\"\n",
        "    def __init__(self, vocab_size, embedding_dim, hidden_size, num_layers, init_range=0.1, bias=True):\n",
        "        super().__init__()\n",
        "        self.vocab_size = vocab_size\n",
        "        self.embedding_dim = embedding_dim\n",
        "        self.hidden_size = hidden_size\n",
        "        self.num_layers = num_layers\n",
        "\n",
        "        self.embedding = Embedding(vocab_size, embedding_dim, init_range)\n",
        "        self.lstm = LSTMLayer(embedding_dim, hidden_size, num_layers, init_range)\n",
        "        self.output_layer = Linear(hidden_size, vocab_size, bias, init_range)\n",
        "\n",
        "        # 存储参数和梯度\n",
        "        self.parameters = self.embedding.parameters + self.lstm.parameters + self.output_layer.parameters\n",
        "        self.gradients = self.embedding.gradients + self.lstm.gradients + self.output_layer.gradients\n",
        "\n",
        "    def forward(self, inputs, h0, c0):\n",
        "        \"\"\"\n",
        "        前向传播\n",
        "\n",
        "        参数:\n",
        "        - inputs: 词语索引序列，形状为 (seq_len, batch_size)\n",
        "        - h0: 初始隐藏状态列表，长度为 num_layers，每个元素形状为 (hidden_size, batch_size)\n",
        "        - c0: 初始细胞状态列表，长度为 num_layers，每个元素形状为 (hidden_size, batch_size)\n",
        "\n",
        "        返回:\n",
        "        - outputs: 每个时间步的输出，形状为 (vocab_size, seq_len * batch_size)\n",
        "        - (h_n, c_n): 最后一个时间步的隐藏状态和细胞状态列表\n",
        "        \"\"\"\n",
        "        self.seq_len, self.batch_size = inputs.shape\n",
        "        embeddings = self.embedding.forward(inputs)\n",
        "        outputs, (h_n, c_n) = self.lstm.forward(embeddings, h0, c0)\n",
        "        outputs = outputs.reshape(self.hidden_size, -1)  # (hidden_size, seq_len * batch_size)\n",
        "        outputs = self.output_layer.forward(outputs)\n",
        "        return outputs, (h_n, c_n)\n",
        "\n",
        "    def backward(self, doutputs, dh_n, dc_n):  \n",
        "        \"\"\"\n",
        "        反向传播\n",
        "\n",
        "        参数:\n",
        "        - doutputs: 每个时间步的输出梯度，形状为 (vocab_size, seq_len * batch_size)\n",
        "        - dh_n: 最后时间步的隐藏状态梯度列表\n",
        "        - dc_n: 最后时间步的细胞状态梯度列表\n",
        "\n",
        "        返回:\n",
        "        - 无需返回输入梯度\n",
        "        \"\"\"\n",
        "        d_lstm_outputs = self.output_layer.backward(doutputs)\n",
        "        d_lstm_outputs = d_lstm_outputs.reshape(self.hidden_size, self.seq_len, self.batch_size)\n",
        "        d_embeddings = self.lstm.backward(d_lstm_outputs, dh_n, dc_n)\n",
        "        if d_embeddings is not None:\n",
        "            self.embedding.backward(d_embeddings)\n",
        "\n",
        "class LanguageModelTrainer:\n",
        "    \"\"\"\n",
        "    语言模型训练器\n",
        "    \"\"\"\n",
        "    def __init__(self, model, loss_func, optimizer):\n",
        "        self.model = model\n",
        "        self.loss_func = loss_func\n",
        "        self.optimizer = optimizer\n",
        "\n",
        "    def train(self, inputs, targets, h0, c0, grad_clip=5):\n",
        "        \"\"\"\n",
        "        训练一个批次的数据\n",
        "\n",
        "        参数:\n",
        "        - inputs: 输入词语索引序列，形状为 (seq_len, batch_size)\n",
        "        - targets: 目标词语索引序列，形状为 (seq_len, batch_size)\n",
        "        - h0: 初始隐藏状态列表\n",
        "        - c0: 初始细胞状态列表\n",
        "        - grad_clip: 梯度裁剪阈值，默认为 5\n",
        "\n",
        "        返回:\n",
        "        - loss: 标量，平均损失值\n",
        "        \"\"\"\n",
        "        self.model.zero_grad()\n",
        "        outputs, (h_n, c_n) = self.model.forward(inputs, h0, c0)\n",
        "        loss = self.loss_func.forward(outputs, targets.flatten())\n",
        "        dloss = self.loss_func.backward()\n",
        "        dh_n = [np.zeros_like(h) for h in h_n]\n",
        "        dc_n = [np.zeros_like(c) for c in c_n]\n",
        "        self.model.backward(dloss, dh_n, dc_n)\n",
        "\n",
        "        # 裁剪梯度\n",
        "        clipped_gradients = [np.clip(grad, -grad_clip, grad_clip) for grad in self.model.gradients]\n",
        "\n",
        "        # 应用梯度\n",
        "        self.optimizer.step(clipped_gradients)\n",
        "        return loss\n",
        "\n",
        "def generate_dummy_data():\n",
        "    # 为了更好地模拟实际文本，包含空格分隔的中文词语和标点符号\n",
        "    text = '你好世界。 机器学习是有趣的。 你好机器。 '\n",
        "    return text\n",
        "\n",
        "def train_model():\n",
        "    # 生成数据\n",
        "    text = generate_dummy_data()\n",
        "    tokenizer = Tokenizer(text)\n",
        "    data = tokenizer.text_to_sequence(text)\n",
        "\n",
        "    # 初始化模型和优化器\n",
        "    vocab_size = tokenizer.vocab_size\n",
        "    embedding_dim = 10\n",
        "    hidden_size = 20\n",
        "    num_layers = 1\n",
        "    model = LanguageModel(vocab_size, embedding_dim, hidden_size, num_layers)\n",
        "    loss_func = CrossEntropyLoss()\n",
        "    optimizer = AdamOptimizer(model.parameters)\n",
        "    trainer = LanguageModelTrainer(model, loss_func, optimizer)\n",
        "\n",
        "    # 准备数据\n",
        "    seq_len = len(data) - 1\n",
        "    batch_size = 1\n",
        "    inputs = np.array([data[:-1]]).T  # (seq_len, batch_size)\n",
        "    targets = np.array([data[1:]]).T  # (seq_len, batch_size)\n",
        "    h0 = [np.zeros((hidden_size, batch_size)) for _ in range(num_layers)]\n",
        "    c0 = [np.zeros((hidden_size, batch_size)) for _ in range(num_layers)]\n",
        "\n",
        "    # 训练\n",
        "    epochs = 1000\n",
        "    for epoch in range(epochs):\n",
        "        loss = trainer.train(inputs, targets, h0, c0)\n",
        "        if (epoch + 1) % 100 == 0:\n",
        "            print(f'Epoch {epoch+1}, Loss: {loss:.4f}')\n",
        "\n",
        "    return model, tokenizer\n",
        "\n",
        "def generate_text(model, tokenizer, seed_text, max_len=20, temperature=1.0):\n",
        "    generated_text = seed_text\n",
        "    h = [np.zeros((model.hidden_size, 1)) for _ in range(model.num_layers)]\n",
        "    c = [np.zeros((model.hidden_size, 1)) for _ in range(model.num_layers)]\n",
        "\n",
        "    # 使用种子文本初始化隐藏状态和细胞状态\n",
        "    seed_sequence = tokenizer.text_to_sequence(seed_text)\n",
        "    for word_id in seed_sequence[:-1]:\n",
        "        inputs = np.array([[word_id]])  # (seq_len=1, batch_size=1)\n",
        "        _, (h, c) = model.forward(inputs, h, c)\n",
        "\n",
        "    last_word_id = seed_sequence[-1]\n",
        "\n",
        "    for _ in range(max_len):\n",
        "        inputs = np.array([[last_word_id]])  # (seq_len=1, batch_size=1)\n",
        "        outputs, (h, c) = model.forward(inputs, h, c)\n",
        "        logits = outputs[:, -1]  # (vocab_size,)\n",
        "        logits = logits / temperature\n",
        "        probs = np.exp(logits - np.max(logits))  # 防止溢出\n",
        "        probs /= np.sum(probs)\n",
        "        next_word_id = np.random.choice(len(probs), p=probs)\n",
        "        next_word = tokenizer.ix_to_word[next_word_id]\n",
        "        if next_word == '<EOS>':\n",
        "            break\n",
        "        generated_text += next_word\n",
        "        last_word_id = next_word_id\n",
        "    return generated_text\n",
        "\n",
        "# 主程序\n",
        "if __name__ == \"__main__\":\n",
        "    # 训练模型\n",
        "    model, tokenizer = train_model()\n",
        "\n",
        "    # 示例生成\n",
        "    seed_text = '你好'\n",
        "    generated = generate_text(model, tokenizer, seed_text)\n",
        "    print(f'Generated Text: {generated}')\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "numpy_lstm (4).ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
